package com.eonmind.ai.core.chat;

import cn.hutool.core.net.url.UrlBuilder;
import com.alibaba.cloud.ai.advisor.DocumentRetrievalAdvisor;
import com.alibaba.cloud.ai.dashscope.api.DashScopeApi;
import com.alibaba.cloud.ai.dashscope.rag.DashScopeDocumentRetriever;
import com.alibaba.cloud.ai.dashscope.rag.DashScopeDocumentRetrieverOptions;
import com.eonmind.ai.core.chat.advisors.token.TokenCollect;
import com.eonmind.ai.core.chat.manager.ChatModelManager;
import com.eonmind.ai.core.chat.model.AiChatModel;
import com.eonmind.ai.mapper.AiAssistantMapper;
import com.eonmind.common.provider.entity.ai.AiAssistant;
import com.eonmind.knowledge.contract.api.KnowledgeFeignService;
import jakarta.annotation.Resource;
import lombok.AllArgsConstructor;
import lombok.Data;
import lombok.NoArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.ObjectUtils;
import org.apache.commons.lang3.StringUtils;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.api.Advisor;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.model.Generation;
import org.springframework.ai.chat.prompt.DefaultChatOptions;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.model.Media;
import org.springframework.ai.rag.retrieval.search.DocumentRetriever;
import org.springframework.stereotype.Service;
import org.springframework.util.MimeTypeUtils;
import reactor.core.publisher.Flux;

import java.net.URL;
import java.util.ArrayList;
import java.util.List;
import java.util.Optional;

import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY;
import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_RETRIEVE_SIZE_KEY;

/**
 * @author Xujie
 * @since 2025/5/5 19:47
 **/
@Slf4j
@Service
public class AiChatService {
    @Resource
    private ChatModelManager chatModelManager;

    @Resource
    private MessageChatMemoryAdvisor mongoChatMemoryAdvisor;

    @Resource
    private TokenCollect tokenCollect;

    @Resource
    private AiAssistantMapper aiAssistantMapper;

    @Resource
    private DashScopeApi dashScopeApi;

    @Resource
    private KnowledgeFeignService knowledgeFeignService;

    public Flux<AssistantMessage> chat(AiChatModel aiChatModel) {
        DefaultChatOptions defaultChatOptions = new DefaultChatOptions();
        defaultChatOptions.setModel(aiChatModel.getModel());

        // 构造用户信息与文件
        Message userMessage = null;
        if (ObjectUtils.isNotEmpty(aiChatModel.getFileList())) {
            List<Media> list = aiChatModel.getFileList().stream()
                    .map(f -> {
                        URL url = UrlBuilder.of(f)
                                .toURL();
                        return new Media(MimeTypeUtils.IMAGE_PNG, url);
                    }).toList();
            userMessage = new UserMessage(aiChatModel.getMessage(), list);
        } else {
            userMessage = new UserMessage(aiChatModel.getMessage());
        }
        Prompt prompt = new Prompt(userMessage, defaultChatOptions);
        ChatModel chatModel = chatModelManager.getChatModel(aiChatModel.getModelId());

        // 构造assistant
        String customSystemPrompt = "";
        if (StringUtils.isNotBlank(aiChatModel.getCustomSystemPrompt())) {
            customSystemPrompt = aiChatModel.getCustomSystemPrompt();
        } else {
            AiAssistant aiAssistant = aiAssistantMapper.selectById(aiChatModel.getAssistantId());
            if (aiAssistant != null) {
                customSystemPrompt = aiAssistant.getAssistantPrompt();
            }
        }
        // 处理知识库
        List<Advisor> advisors = new ArrayList<>(List.of(mongoChatMemoryAdvisor, tokenCollect));
        if (StringUtils.isNotBlank(aiChatModel.getKnowledgeIndexName())) {
            DocumentRetriever retriever = new DashScopeDocumentRetriever(dashScopeApi,
                    DashScopeDocumentRetrieverOptions.builder().withIndexName(aiChatModel.getKnowledgeIndexName()).build());
            advisors.add(new DocumentRetrievalAdvisor(retriever));
        }
        Advisor[] advisorsArray = advisors.toArray(Advisor[]::new);
        ChatClient chatClient = buildChatClientWithMemory(chatModel
                , aiChatModel.getConversationId()
                , customSystemPrompt
                , advisorsArray);
        return chatClient
                .prompt(prompt).stream().chatResponse()
                .map(
                        t -> {
                            log.info("AI的响应：{}", t);
                            return Optional.ofNullable(t)
                                    .map(ChatResponse::getResult)
                                    .map(Generation::getOutput)
                                    .orElse(new AssistantMessage(""));
                        }
                )
                .doOnCancel(() -> {
                    log.info("取消了会话");
                });
    }


    private ChatClient buildChatClientWithMemory(ChatModel chatModel, Long conversationId, String customPrompt, Advisor... advisors) {

        AiCustomSystemPrompt aiCustomSystemPrompt = new AiCustomSystemPrompt(customPrompt);
        ChatClient chatClient = ChatClient
                .builder(chatModel)
                .defaultAdvisors(advisors)
                .defaultSystem(buildSystemPrompt(aiCustomSystemPrompt))
                .defaultAdvisors(
                        a -> a
                                .param(CHAT_MEMORY_CONVERSATION_ID_KEY, conversationId)
                                .param(CHAT_MEMORY_RETRIEVE_SIZE_KEY, 10)
                                .param("Token", 0)
                )
                .build();
        return chatClient;
    }

    private String buildSystemPrompt(AiCustomSystemPrompt aiCustomSystemPrompt) {

        return defaultSystemPrompt + "\n" + aiCustomSystemPrompt.getCustomPrompt();
    }

    @Data
    @NoArgsConstructor
    @AllArgsConstructor
    private static class AiCustomSystemPrompt {
        private String customPrompt = "";
    }


    private final String defaultSystemPrompt = """
             你是一个 Markdown 格式输出专家。你的所有回答必须严格遵守标准 Markdown 语法。具体要求如下：
             1. 所有回答都应使用 Markdown 结构化格式进行排版，如使用标题（#）、列表（- 或 1.）、代码块（```）、引用块（>）等。
             2. 若回答中包含代码，必须使用三重反引号（```）包裹，并指明语言（如 ```python、```json 等）。
             3. 必要时使用加粗（**）和斜体（*）来突出重点。不要使用 HTML 标签。
             4. 出现术语或关键词时建议使用`代码样式`强调。
             5. 数学公式用标准markdown格式返回如:
             $$
             \\begin{aligned}
             f(x) &= x^2 + 2x + 1 \\\\
                  &= (x + 1)^2
             \\end{aligned}
             $$
            
             示例输出格式：
            
             # 标题示例
            
             这是对问题的回答内容。
            
             ## 二级标题
            
             - 要点一
             - 要点二
             - 要点三
            
             ### 示例代码
             ```python
             def hello_world():
                 print("Hello, Markdown!")
             ```
             > 提示：请始终确保返回结果为合法且渲染正确的 Markdown 文本格式。
             从现在开始，请始终以 Markdown 格式进行所有内容输出。
            
            """;

}
